package middleware

import (
	encrypt "api/pkg/encrypt"
	"bytes"
	"fmt"
	"github.com/zeromicro/go-zero/core/logx"
	"github.com/zeromicro/go-zero/rest/httpx"
	"io"
	"net/http"
)

var secret = []byte("03d2801301c642f49e1ebbe96eb3586b")

// 自定义响应写入器，用于捕获响应内容
type loggingResponseWriter struct {
	http.ResponseWriter
	bodyBuf    bytes.Buffer
	statusCode int
}

func (lrw *loggingResponseWriter) WriteHeader(code int) {
	lrw.statusCode = code
	lrw.ResponseWriter.WriteHeader(code)
}

func (lrw *loggingResponseWriter) Write(body []byte) (int, error) {
	// 记录响应内容
	lrw.bodyBuf.Write(body)
	// 这里不实际写入，由中间件统一处理
	return len(body), nil
}

// EncryptConfig 加密配置
type EncryptConfig struct {
	SecretKey string // AES密钥（16/24/32字节分别对应AES-128/192/256）
	Enable    bool   // 是否启用加密
}

// EncryptMiddleware 加密中间件
type EncryptMiddleware struct {
	config EncryptConfig
}

// NewEncryptMiddleware 创建加密中间件实例
func NewEncryptMiddleware(config EncryptConfig) *EncryptMiddleware {
	return &EncryptMiddleware{
		config: config,
	}
}

// Handle 中间件处理函数
func (m *EncryptMiddleware) Handle(next http.HandlerFunc) http.HandlerFunc {
	return func(w http.ResponseWriter, r *http.Request) {
		// 如果未启用加密，直接放行
		if !m.config.Enable {
			next(w, r)
			return
		}

		// 1. 解密请求体
		if r.Body != nil && r.ContentLength > 0 {
			bodyBytes, err := io.ReadAll(r.Body)
			if err != nil {
				httpx.Error(w, fmt.Errorf("读取请求体失败: %v", err))
				return
			}
			defer r.Body.Close()

			// 解密请求体
			decryptedBody, err := encrypt.AESDecrypt(bodyBytes, secret)
			if err != nil {
				httpx.Error(w, fmt.Errorf("解密请求失败: %v", err))
				return
			}

			// 替换请求体
			r.Body = io.NopCloser(bytes.NewBuffer(decryptedBody))
		}

		// 2. 创建自定义响应写入器，用于拦截响应
		lrw := &loggingResponseWriter{ResponseWriter: w, statusCode: http.StatusOK}

		// 3. 继续处理请求
		next(lrw, r)

		// 4. 加密响应体
		if lrw.bodyBuf.Len() > 0 {
			encryptedBody, err := encrypt.AESEncrypt(lrw.bodyBuf.Bytes(), secret)
			if err != nil {
				logx.Errorf("加密响应失败: %v", err)
				return
			}

			// 清空原始响应头
			w.Header().Del("Content-Length")
			w.Header().Set("Content-Type", "application/octet-stream")

			// 写入加密后的响应体
			_, err = w.Write(encryptedBody)
			if err != nil {
				logx.Errorf("写入响应失败: %v", err)
			}
		}
	}
}
